package ink.metoo.gude.aspect

import com.github.benmanes.caffeine.cache.Caffeine
import ink.metoo.gude.handler.ApiException
import ink.metoo.gude.model.ApiBody
import ink.metoo.gude.model.ApiStatus
import ink.metoo.gude.util.IpUtil
import ink.metoo.gude.util.logger
import jakarta.servlet.http.HttpServletRequest
import org.aspectj.lang.ProceedingJoinPoint
import org.aspectj.lang.annotation.Around
import org.aspectj.lang.annotation.Aspect
import org.springframework.cache.Cache
import org.springframework.cache.caffeine.CaffeineCacheManager
import org.springframework.cache.set
import org.springframework.security.core.context.SecurityContextHolder
import java.util.*
import java.util.concurrent.ConcurrentHashMap

@Aspect
class RateLimiterAspect(
    private val caffeineCacheManager: CaffeineCacheManager,
    private val request: HttpServletRequest
) {

    private val rateLimitIdMap = ConcurrentHashMap<RateLimit, String>()
    private val registerCustomCacheNames = mutableListOf<String>()
    private val log = logger()

    @Around("@within(annotation) || @annotation(annotation)")
    fun pointcut(point: ProceedingJoinPoint, annotation: RateLimit): Any? {
        val key = rateLimitIdMap.compute(annotation) { _, value ->
            value ?: UUID.randomUUID().toString()
        }
        val cacheKey = "rate-limit:${key}"
        if (!registerCustomCacheNames.contains(cacheKey)) {
            log.info("Registering custom cache for $cacheKey")
            caffeineCacheManager.registerCustomCache(
                cacheKey, Caffeine.newBuilder()
                    .maximumSize(Int.MAX_VALUE.toLong())
                    .expireAfterWrite(annotation.duration, annotation.timeUnit)
                    .recordStats()
                    .build()
            )
            registerCustomCacheNames.add(cacheKey)
        }
        val cache = caffeineCacheManager.getCache(cacheKey)!!
        val name = when (annotation.key) {
            RateLimit.Key.USER -> {
                val authentication =
                    SecurityContextHolder.getContext().authentication ?: throw ApiException(ApiStatus.U_2003)
                authentication.name
            }

            RateLimit.Key.IP -> {
                IpUtil.getClientIp(request)
            }
        }
        cache.check(name, annotation.max)
        try {
            val proceed = point.proceed()
            if (annotation.type == RateLimit.Type.RESPONSE_SUCCESS) {
                if (proceed !is ApiBody<*> || proceed.code == ApiStatus.SUCCESS.code) {
                    cache.increment(name)
                }
            } else {
                cache.increment(name)
            }
            return proceed
        } catch (e: Exception) {
            if (annotation.type == RateLimit.Type.ALL) {
                cache.increment(name)
            }
            throw e
        }
    }

    fun Cache.increment(name: String) {
        if (this[name] == null) {
            this[name] = 1
            return
        }
        val uc = this[name, Int::class.javaObjectType]!!
        this[name] = (uc + 1)
    }

    fun Cache.check(name: String, max: Int) {
        if (this[name] == null) {
            return
        }
        val uc = this[name, Int::class.javaObjectType]!!
        if (uc >= max) {
            throw ApiException(ApiStatus.U_2009, mapOf("key" to name))
        }
    }

}